#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
四足机器人轨迹可视化模块
包含末端执行器标记和轨迹可视化功能
"""

import pybullet as p
import numpy as np
from config import (
    ENABLE_END_EFFECTOR_MARKERS, MARKER_COLOR, MARKER_RADIUS,
    ENABLE_TRAJECTORY_VIZ, MAX_TRAJ_SEGMENTS, TRAJECTORY_LINE_COLOR, TRAJECTORY_LINE_WIDTH
)

class TrajectoryVisualizer:
    """轨迹可视化类，负责末端执行器标记和轨迹可视化"""
    
    def __init__(self):
        """初始化轨迹可视化器"""
        self.special_markers = {}
        self.trajectory_viz_points = {}
        self.trajectory_line_ids = {}
    
    def create_end_effector_markers(self, leg_ids):
        """为每条腿创建末端执行器标记
        
        参数:
            leg_ids: 腿的ID列表
            
        返回:
            special_markers: 末端执行器标记字典
        """
        if not ENABLE_END_EFFECTOR_MARKERS:
            return {}
            
        for leg_id in leg_ids:
            # 创建红色球体可视化目标点
            special_marker_visual = p.createVisualShape(
                p.GEOM_SPHERE, 
                radius=MARKER_RADIUS, 
                rgbaColor=MARKER_COLOR
            )
            
            # 初始位置设置在远处
            special_marker_id = p.createMultiBody(
                baseVisualShapeIndex=special_marker_visual,
                basePosition=[100, 100, 100]
            )
            
            self.special_markers[leg_id] = special_marker_id
            self.trajectory_viz_points[leg_id] = []
            self.trajectory_line_ids[leg_id] = []
        
        return self.special_markers
    
    def update_marker_position(self, leg_id, target_pos_world):
        """更新标记位置
        
        参数:
            leg_id: 腿的ID
            target_pos_world: 目标位置 (世界坐标系)
        """
        if not ENABLE_END_EFFECTOR_MARKERS or leg_id not in self.special_markers:
            return
            
        # 更新标记位置
        p.resetBasePositionAndOrientation(
            self.special_markers[leg_id], 
            target_pos_world, 
            [0, 0, 0, 1]
        )
        
        # 更新轨迹可视化
        if not ENABLE_TRAJECTORY_VIZ:
            return
            
        if len(self.trajectory_viz_points[leg_id]) > 0:
            prev_ee_world_pos = self.trajectory_viz_points[leg_id][-1]
            line_id = p.addUserDebugLine(
                prev_ee_world_pos, 
                target_pos_world,
                lineColorRGB=TRAJECTORY_LINE_COLOR,
                lineWidth=TRAJECTORY_LINE_WIDTH
            )
            self.trajectory_line_ids[leg_id].append(line_id)
        
        self.trajectory_viz_points[leg_id].append(target_pos_world)
        
        # 限制轨迹段数量
        if len(self.trajectory_line_ids[leg_id]) > MAX_TRAJ_SEGMENTS:
            p.removeUserDebugItem(self.trajectory_line_ids[leg_id].pop(0))
            self.trajectory_viz_points[leg_id].pop(0)
    
    def clear_trajectories(self):
        """清除所有轨迹"""
        if not ENABLE_TRAJECTORY_VIZ:
            return
            
        for leg_id in self.trajectory_line_ids:
            for line_id in self.trajectory_line_ids[leg_id]:
                p.removeUserDebugItem(line_id)
            self.trajectory_line_ids[leg_id] = []
            self.trajectory_viz_points[leg_id] = []
